import os
import torch
import cv2 as cv
import pandas as pd
from torch.utils.data import Dataset


# 定义一个dataset
class MyDataset(Dataset):
    def __init__(self, dataset_folder, transform=None):
        self.dataset_folder = dataset_folder  # k线图文件夹
        self.label_data_name = os.path.join(dataset_folder,  dataset_folder+".csv")  # 标签数据文件
        label_data = pd.read_csv(self.label_data_name, index_col=0)  # 读取标签数据,设置第0列为index
        self.y1 = label_data.iloc[:]['filename'].values  # 取第一列,图片名
        self.y3 = torch.from_numpy(label_data.iloc[:]['label'].values)  # 取第三列,并转化为tensor
        self.len = len(label_data)
        self.transform = transform

    def __getitem__(self, index):
        img_name = self.y1[index]
        img_path = os.path.join(self.dataset_folder, img_name)  # 文件夹路径组合
        image = cv.imread(img_path)  # 读取图片，[287, 287, 3]
        if self.transform is not None:  # 类型转化，转化为tensor，并正则化
            img = self.transform(image)
        return img, self.y3[index]

    def __len__(self):
        return self.len


'''https://www.pytorchtutorial.com/pytorch-custom-dataset-examples/'''
